/**
 * ============================================================================
 *
 * Copyright (C) 2018, Hisilicon Technologies Co., Ltd. All Rights Reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 *   1 Redistributions of source code must retain the above copyright notice,
 *     this list of conditions and the following disclaimer.
 *
 *   2 Redistributions in binary form must reproduce the above copyright notice,
 *     this list of conditions and the following disclaimer in the documentation
 *     and/or other materials provided with the distribution.
 *
 *   3 Neither the names of the copyright holders nor the names of the
 *   contributors may be used to endorse or promote products derived from this
 *   software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 * ============================================================================
 */
#include <vector>
#include <sstream>
#include <cmath>
#include <sys/time.h>
#include <opencv2/opencv.hpp>
#include "hiaiengine/log.h"
#include "ascenddk/ascend_ezdvpp/dvpp_data_type.h"
#include "ascenddk/ascend_ezdvpp/dvpp_process.h"
#include "i420_compress_jpg.h"
#include "util.hpp"

using hiai::Engine;
using namespace std::__cxx11;

// register data type
HIAI_REGISTER_DATA_TYPE("EngineTransT", EngineTransT);
HIAI_REGISTER_DATA_TYPE("OutputT", OutputT);
HIAI_REGISTER_DATA_TYPE("NewImageParaT", NewImageParaT);

// constants
namespace {
// output port (engine port begin with 0)
const uint32_t kSendDataPort = 0;

// call dvpp success
const uint32_t kDvppProcSuccess = 0;
// level for call DVPP
const int32_t kDvppToJpegLevel = 100;
}

I420CompressJpg::I420CompressJpg() {
  compress_config_ = nullptr;
}

HIAI_StatusT I420CompressJpg::Init(
    const hiai::AIConfig& config,
    const std::vector<hiai::AIModelDescription>& model_desc) {
  HIAI_ENGINE_LOG("Begin initialize!");

  // get configurations
  if (compress_config_ == nullptr) {
    compress_config_ = std::make_shared<ResizedImageConfig>();
  }

  // get parameters from graph.config
  for (int index = 0; index < config.items_size(); index++) {
    const ::hiai::AIConfigItem& item = config.items(index);
    const std::string& name = item.name();
    const std::string& value = item.value();

    if (name == "resized_width") {
      compress_config_->resized_width = atoi(value.data());
    } else if (name == "resized_height") {
      compress_config_->resized_height = atoi(value.data());
    } else {
      HIAI_ENGINE_LOG("unused config name: %s", name.c_str());
    }
  }

  HIAI_ENGINE_LOG(HIAI_DEBUG_INFO, "End initialize!");
  return HIAI_OK;
}

HIAI_StatusT I420CompressJpg::SendResults(
    const std::shared_ptr<EngineTransT> &inference_res) {
  HIAI_StatusT status = HIAI_OK;
  std::shared_ptr<EngineTransT> trans_data = std::make_shared<EngineTransT>();
  trans_data->b_info = inference_res->b_info;
  trans_data->status = true;

  std::vector<OutputT> output_data_vec = inference_res->output_datas;
  // dealing every image
  std::vector<NewImageParaT> converted_jpeg;
  for (uint32_t ind = 0; ind < inference_res->b_info.batch_size; ind++) {
    OutputT out = output_data_vec[ind];
    int out_size = out.size;
    std::shared_ptr<u_int8_t> out_data(new u_int8_t[out_size]);
    memcpy_s(out_data.get(), out_size, out.data.get(), out_size);

    //**** convert I420 image to NV12, then call encapsulated api to compress nv12 to jpeg.
    unsigned long convert_start_time = get_current_time();
    NewImageParaT new_image;
    uint32_t len_y = out_size * 2 / 3;
    uint32_t len_u = len_y/4;
    new_image.img.channel = 0;
    new_image.img.format = hiai::YUV420SP;
    new_image.img.width = compress_config_->resized_width;
    new_image.img.height = compress_config_->resized_height;
    new_image.img.size = out_size;
    std::shared_ptr<u_int8_t> share_data(new u_int8_t[out_size]);
    memset_s(share_data.get(), out_size, 0, out_size);
    if (memcpy_s(share_data.get(), len_y, out_data.get(), len_y)) {
      HIAI_ENGINE_LOG(HIAI_ENGINE_RUN_ARGS_NOT_RIGHT, "memcpy failed.");
      return HIAI_ERROR;
    }
    // Process uv component.
    u_int8_t* ptr = share_data.get();
    for (int i = 0; i < len_u; ++i) {
      ptr[len_y + 2*i] = out_data.get()[len_y + i];
      ptr[len_y + 2*i+1] = out_data.get()[len_y + len_u + i];
    }
    new_image.img.data = share_data;
    HIAI_StatusT convert_ret = ConvertImage(new_image);
    if (convert_ret != HIAI_OK) {
      HIAI_ENGINE_LOG(HIAI_ENGINE_RUN_ARGS_NOT_RIGHT, "Convert YUV Image to Jpeg failed!");
      return HIAI_ERROR;
    }
    unsigned long convert_time = get_current_time() - convert_start_time;
    HIAI_ENGINE_LOG(HIAI_ENGINE_RUN_ARGS_NOT_RIGHT, "convert to jpg used time:%d (us)", convert_time);
    converted_jpeg.push_back(new_image);
    trans_data->imgs = converted_jpeg;
  }
  // send results and original image data to post process (port 0)
  HIAI_StatusT hiai_ret = SendData(kSendDataPort, "EngineTransT",
                                   std::static_pointer_cast<void>(trans_data));
  
  return hiai_ret;
}

bool I420CompressJpg::IsSupportFormat(hiai::IMAGEFORMAT format) {
  return format == hiai::YUV420SP;
}

HIAI_StatusT I420CompressJpg::ConvertImage(NewImageParaT& img) {
  hiai::IMAGEFORMAT format = img.img.format;
  if (!IsSupportFormat(format)){
    HIAI_ENGINE_LOG(HIAI_ENGINE_RUN_ARGS_NOT_RIGHT,
                    "Format %d is not supported!", format);
    return HIAI_ERROR;
  }

  uint32_t width = img.img.width;
  uint32_t height = img.img.height;
  uint32_t img_size = img.img.size;
  // HIAI_ENGINE_LOG(HIAI_ENGINE_RUN_ARGS_NOT_RIGHT, "ConvertImage w:%d h:%d size:%d.", width, height, img_size);

  // parameter
  ascend::utils::DvppToJpgPara dvpp_to_jpeg_para;
  dvpp_to_jpeg_para.format = JPGENC_FORMAT_NV12;
  dvpp_to_jpeg_para.level = kDvppToJpegLevel;
  dvpp_to_jpeg_para.resolution.height = height;
  dvpp_to_jpeg_para.resolution.width = width;
  ascend::utils::DvppProcess dvpp_to_jpeg(dvpp_to_jpeg_para);

  // call DVPP
  ascend::utils::DvppOutput dvpp_output;
  int32_t ret = dvpp_to_jpeg.DvppOperationProc(reinterpret_cast<char*>(img.img.data.get()),
                                                img_size, &dvpp_output);

  // failed, no need to send to presenter
  if (ret != kDvppProcSuccess) {
    HIAI_ENGINE_LOG(HIAI_ENGINE_RUN_ARGS_NOT_RIGHT,
                    "Failed to convert YUV420SP to JPEG, skip it.");
    return HIAI_ERROR;
  }

  // reset the data in img_vec
  img.img.data.reset(dvpp_output.buffer, default_delete<uint8_t[]>());
  img.img.size = dvpp_output.size;

  return HIAI_OK;
}

HIAI_IMPL_ENGINE_PROCESS("i420_compress_jpg",
    I420CompressJpg, INPUT_SIZE) {
  // check arg0 is null or not
  if (arg0 == nullptr) {
    HIAI_ENGINE_LOG(HIAI_ENGINE_RUN_ARGS_NOT_RIGHT,
                    "Failed to process invalid message.");
    return HIAI_ERROR;
  }

  // check original image is empty or not
  std::shared_ptr<EngineTransT> inference_res = std::static_pointer_cast<
      EngineTransT>(arg0);

  // dealing inference results
  return SendResults(inference_res);
}
